#include "Parser.h"
#include <stdexcept>
#include <iostream>

namespace MiniDB
{

    Parser::Parser(const std::vector<Token> &tokens)
        : tokens_(tokens), position_(0), current_token_(TokenType::EOF_TOKEN, "", 0, 0)
    {
        if (!tokens_.empty())
        {
            current_token_ = tokens_[0];
        }
    }

    void Parser::advance()
    {
        position_++;
        if (position_ >= tokens_.size())
        {
            current_token_ = Token(TokenType::EOF_TOKEN, "", 0, 0);
        }
        else
        {
            current_token_ = tokens_[position_];
        }
    }

    Token Parser::peek() const
    {
        size_t peek_pos = position_ + 1;
        if (peek_pos >= tokens_.size())
        {
            return Token(TokenType::EOF_TOKEN, "", 0, 0);
        }
        return tokens_[peek_pos];
    }

    Token Parser::expect(TokenType type)
    {
        if (current_token_.type != type)
        {
            std::string expected = "EXPECTED_TYPE";  // TODO: Map type to string
            std::string actual = std::to_string(static_cast<int>(current_token_.type));
            throw std::runtime_error("Expected " + expected + ", but got token " + actual +
                                     " at line " + std::to_string(current_token_.line));
        }
        Token token = current_token_;
        advance();
        return token;
    }

    std::vector<std::shared_ptr<Statement>> Parser::parse()
    {
        std::vector<std::shared_ptr<Statement>> statements;

        while (current_token_.type != TokenType::EOF_TOKEN)
        {
            auto statement = parseStatement();
            statements.push_back(statement);

            // 跳过分号
            if (current_token_.type == TokenType::SEMICOLON)
            {
                advance();
            }
        }

        return statements;
    }

    std::shared_ptr<Statement> Parser::parseStatement()
    {
        if (current_token_.type == TokenType::SELECT)
        {
            return parseSelectStatement();
        }
        else if (current_token_.type == TokenType::INSERT)
        {
            return parseInsertStatement();
        }
        else if (current_token_.type == TokenType::UPDATE)
        {
            return parseUpdateStatement();
        }
        else if (current_token_.type == TokenType::DELETE)
        {
            return parseDeleteStatement();
        }
        else if (current_token_.type == TokenType::CREATE)
        {
            return parseCreateStatement();
        }
        else if (current_token_.type == TokenType::DROP)
        {
            return parseDropStatement();
        }
        else
        {
            throw std::runtime_error("Unexpected token: " + std::to_string(static_cast<int>(current_token_.type)));
        }
    }

    std::shared_ptr<SelectStatement> Parser::parseSelectStatement()
    {
        expect(TokenType::SELECT);

        // 解析选择列表
        auto select_list = parseSelectList();

        // 解析FROM子句
        expect(TokenType::FROM);
        auto from_clause = parseFromClause();

        // 可选的WHERE子句
        std::shared_ptr<Expression> where_clause = nullptr;
        if (current_token_.type == TokenType::WHERE)
        {
            advance();
            where_clause = parseExpression();
        }

        return std::make_shared<SelectStatement>(select_list, from_clause, where_clause);
    }

    std::vector<std::shared_ptr<Expression>> Parser::parseSelectList()
    {
        std::vector<std::shared_ptr<Expression>> select_list;

        if (current_token_.type == TokenType::STAR)
        {
            select_list.push_back(std::make_shared<ColumnReference>("*"));
            advance();
        }
        else
        {
            select_list.push_back(parseExpression());
            while (current_token_.type == TokenType::COMMA)
            {
                advance();
                select_list.push_back(parseExpression());
            }
        }

        return select_list;
    }

    std::shared_ptr<FromClause> Parser::parseFromClause()
    {
        std::string table_name = expect(TokenType::IDENTIFIER).value;

        std::string alias;
        if (current_token_.type == TokenType::IDENTIFIER)
        {
            alias = current_token_.value;
            advance();
        }

        return std::make_shared<FromClause>(table_name, alias);
    }

    std::shared_ptr<InsertStatement> Parser::parseInsertStatement()
    {
        expect(TokenType::INSERT);
        expect(TokenType::INTO);

        std::string table_name = expect(TokenType::IDENTIFIER).value;

        // 可选的列列表
        std::vector<std::string> columns;
        if (current_token_.type == TokenType::LPAREN)
        {
            advance();
            columns.push_back(expect(TokenType::IDENTIFIER).value);
            while (current_token_.type == TokenType::COMMA)
            {
                advance();
                columns.push_back(expect(TokenType::IDENTIFIER).value);
            }
            expect(TokenType::RPAREN);
        }

        expect(TokenType::VALUES);

        // 解析值列表
        std::vector<std::vector<std::shared_ptr<Expression>>> values;
        expect(TokenType::LPAREN);
        std::vector<std::shared_ptr<Expression>> row_values;
        row_values.push_back(parseExpression());
        while (current_token_.type == TokenType::COMMA)
        {
            advance();
            row_values.push_back(parseExpression());
        }
        expect(TokenType::RPAREN);
        values.push_back(row_values);

        return std::make_shared<InsertStatement>(table_name, columns, values);
    }

    std::shared_ptr<UpdateStatement> Parser::parseUpdateStatement()
    {
        expect(TokenType::UPDATE);

        std::string table_name = expect(TokenType::IDENTIFIER).value;

        expect(TokenType::SET);

        // 解析赋值列表
        std::vector<std::pair<std::string, std::shared_ptr<Expression>>> assignments;
        std::string column_name = expect(TokenType::IDENTIFIER).value;
        expect(TokenType::EQUALS);
        auto expression = parseExpression();
        assignments.emplace_back(column_name, expression);

        while (current_token_.type == TokenType::COMMA)
        {
            advance();
            column_name = expect(TokenType::IDENTIFIER).value;
            expect(TokenType::EQUALS);
            expression = parseExpression();
            assignments.emplace_back(column_name, expression);
        }

        // 可选的WHERE子句
        std::shared_ptr<Expression> where_clause = nullptr;
        if (current_token_.type == TokenType::WHERE)
        {
            advance();
            where_clause = parseExpression();
        }

        return std::make_shared<UpdateStatement>(table_name, assignments, where_clause);
    }

    std::shared_ptr<DeleteStatement> Parser::parseDeleteStatement()
    {
        expect(TokenType::DELETE);
        expect(TokenType::FROM);

        std::string table_name = expect(TokenType::IDENTIFIER).value;

        // 可选的WHERE子句
        std::shared_ptr<Expression> where_clause = nullptr;
        if (current_token_.type == TokenType::WHERE)
        {
            advance();
            where_clause = parseExpression();
        }

        return std::make_shared<DeleteStatement>(table_name, where_clause);
    }

    std::shared_ptr<CreateTableStatement> Parser::parseCreateStatement()
    {
        expect(TokenType::CREATE);
        expect(TokenType::TABLE);

        std::string table_name = expect(TokenType::IDENTIFIER).value;

        expect(TokenType::LPAREN);

        // 解析列定义
        std::vector<std::shared_ptr<ASTNode>> columns;
        columns.push_back(parseColumnDefinition());

        while (current_token_.type == TokenType::COMMA)
        {
            advance();
            columns.push_back(parseColumnDefinition());
        }

        expect(TokenType::RPAREN);

        return std::make_shared<CreateTableStatement>(table_name, columns);
    }

    std::shared_ptr<DropTableStatement> Parser::parseDropStatement()
    {
        expect(TokenType::DROP);
        expect(TokenType::TABLE);

        std::string table_name = expect(TokenType::IDENTIFIER).value;

        return std::make_shared<DropTableStatement>(table_name);
    }

    std::shared_ptr<ColumnDefinition> Parser::parseColumnDefinition()
    {
        std::string column_name = expect(TokenType::IDENTIFIER).value;

        // 解析数据类型
        std::string data_type = parseDataType();

        // 可选的约束
        std::vector<std::string> constraints;
        while (current_token_.type == TokenType::IDENTIFIER ||
               current_token_.type == TokenType::PRIMARY ||
               current_token_.type == TokenType::KEY ||
               current_token_.type == TokenType::NOT ||
               current_token_.type == TokenType::NULL_TOKEN)
        {
            std::string constraint = current_token_.value;
            if (constraint == "PRIMARY" || constraint == "NOT" || constraint == "UNIQUE" || constraint == "AUTO_INCREMENT" ||
                current_token_.type == TokenType::PRIMARY || current_token_.type == TokenType::NOT)
            {
                constraints.push_back(constraint);
                advance();
                // 处理复合约束如 PRIMARY KEY, NOT NULL
                if (constraint == "PRIMARY" && current_token_.type == TokenType::KEY)
                {
                    constraints.back() = "PRIMARY KEY";
                    advance();
                }
                else if (constraint == "NOT" && current_token_.type == TokenType::NULL_TOKEN)
                {
                    constraints.back() = "NOT NULL";
                    advance();
                }
            }
            else
            {
                break;
            }
        }

        return std::make_shared<ColumnDefinition>(column_name, data_type, constraints);
    }

    std::string Parser::parseDataType()
    {
        if (current_token_.type == TokenType::INT ||
            current_token_.type == TokenType::VARCHAR ||
            current_token_.type == TokenType::FLOAT ||
            current_token_.type == TokenType::BOOLEAN ||
            current_token_.type == TokenType::DATE ||
            current_token_.type == TokenType::TEXT)
        {

            std::string data_type = current_token_.value;
            advance();

            // 处理VARCHAR(n)的情况
            if (data_type == "VARCHAR" && current_token_.type == TokenType::LPAREN)
            {
                advance();
                std::string size = expect(TokenType::NUMBER_LITERAL).value;
                expect(TokenType::RPAREN);
                data_type = "VARCHAR(" + size + ")";
            }

            return data_type;
        }
        else
        {
            throw std::runtime_error("Expected data type, but got " + std::to_string(static_cast<int>(current_token_.type)));
        }
    }

    // 表达式解析方法
    std::shared_ptr<Expression> Parser::parseExpression()
    {
        return parseOrExpression();
    }

    std::shared_ptr<Expression> Parser::parseOrExpression()
    {
        auto left = parseAndExpression();

        while (current_token_.type == TokenType::OR)
        {
            std::string operator_ = current_token_.value;
            advance();
            auto right = parseAndExpression();
            left = std::make_shared<BinaryExpression>(left, operator_, right);
        }

        return left;
    }

    std::shared_ptr<Expression> Parser::parseAndExpression()
    {
        auto left = parseNotExpression();

        while (current_token_.type == TokenType::AND)
        {
            std::string operator_ = current_token_.value;
            advance();
            auto right = parseNotExpression();
            left = std::make_shared<BinaryExpression>(left, operator_, right);
        }

        return left;
    }

    std::shared_ptr<Expression> Parser::parseNotExpression()
    {
        if (current_token_.type == TokenType::NOT)
        {
            advance();
            auto operand = parseComparisonExpression();
            return std::make_shared<UnaryExpression>("NOT", operand);
        }
        else
        {
            return parseComparisonExpression();
        }
    }

    std::shared_ptr<Expression> Parser::parseComparisonExpression()
    {
        auto left = parseAdditiveExpression();

        if (current_token_.type == TokenType::EQUALS ||
            current_token_.type == TokenType::NOT_EQUALS ||
            current_token_.type == TokenType::GREATER ||
            current_token_.type == TokenType::LESS ||
            current_token_.type == TokenType::GREATER_EQUALS ||
            current_token_.type == TokenType::LESS_EQUALS ||
            current_token_.type == TokenType::LIKE)
        {
            std::string operator_ = current_token_.value;
            advance();
            auto right = parseAdditiveExpression();
            return std::make_shared<BinaryExpression>(left, operator_, right);
        }

        return left;
    }

    std::shared_ptr<Expression> Parser::parseAdditiveExpression()
    {
        auto left = parseMultiplicativeExpression();

        while (current_token_.value == "+" || current_token_.value == "-")
        {
            std::string operator_ = current_token_.value;
            advance();
            auto right = parseMultiplicativeExpression();
            left = std::make_shared<BinaryExpression>(left, operator_, right);
        }

        return left;
    }

    std::shared_ptr<Expression> Parser::parseMultiplicativeExpression()
    {
        auto left = parsePrimaryExpression();

        while (current_token_.value == "*" || current_token_.value == "/")
        {
            std::string operator_ = current_token_.value;
            advance();
            auto right = parsePrimaryExpression();
            left = std::make_shared<BinaryExpression>(left, operator_, right);
        }

        return left;
    }

    std::shared_ptr<Expression> Parser::parsePrimaryExpression()
    {
        if (current_token_.type == TokenType::IDENTIFIER)
        {
            std::string identifier = current_token_.value;
            advance();

            // 检查是否是函数调用
            if (current_token_.type == TokenType::LPAREN)
            {
                advance();
                std::vector<std::shared_ptr<Expression>> arguments;
                if (current_token_.type != TokenType::RPAREN)
                {
                    arguments.push_back(parseExpression());
                    while (current_token_.type == TokenType::COMMA)
                    {
                        advance();
                        arguments.push_back(parseExpression());
                    }
                }
                expect(TokenType::RPAREN);
                return std::make_shared<FunctionCall>(identifier, arguments);
            }
            else
            {
                return std::make_shared<ColumnReference>(identifier);
            }
        }
        else if (current_token_.type == TokenType::STRING_LITERAL)
        {
            std::string value = current_token_.value;
            advance();
            return std::make_shared<Literal>(value, "STRING");
        }
        else if (current_token_.type == TokenType::NUMBER_LITERAL)
        {
            std::string value_str = current_token_.value;
            advance();
            // 尝试转换为数字
            try
            {
                if (value_str.find('.') != std::string::npos)
                {
                    double value = std::stod(value_str);
                    return std::make_shared<Literal>(value, "FLOAT");
                }
                else
                {
                    int value = std::stoi(value_str);
                    return std::make_shared<Literal>(value, "INT");
                }
            }
            catch (const std::exception &)
            {
                return std::make_shared<Literal>(value_str, "STRING");
            }
        }
        else if (current_token_.type == TokenType::LPAREN)
        {
            advance();
            auto expression = parseExpression();
            expect(TokenType::RPAREN);
            return expression;
        }
        else
        {
            throw std::runtime_error("Unexpected token in expression: " +
                                     std::to_string(static_cast<int>(current_token_.type)));
        }
    }

} // namespace MiniDB
